/*******************************************************************************
 * Copyright (C) 2018-2020 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 ******************************************************************************/

#include "identify.h"

#include "config.h"
#include "gst/videoanalytics/video_frame.h"
#include "gstgvaidentify.h"

#include <opencv2/imgproc.hpp>
#include <vector>

#define UNUSED(x) (void)x

using namespace std::placeholders;

Identify::Identify(GstGvaIdentify *ovino) {
    masterGstElement_ = ovino;

    // Create gallery
    if (ovino->gallery) {
        double reid_threshold = ovino->cosine_dist_threshold;
        gallery = std::unique_ptr<EmbeddingsGallery>(
            new EmbeddingsGallery(GST_BASE_TRANSFORM_CAST(ovino), ovino->gallery, reid_threshold));
    }
}

Identify::~Identify() {
}

void Identify::IdentifyObjects(GstBuffer *buffer) {
    GVA::VideoFrame video_frame(buffer, masterGstElement_->info);
    std::vector<cv::Mat> embeddings;
    std::vector<GVA::Tensor> tensors;
    auto regions = video_frame.regions();

    // Find embeddings generated by gvaclassify
    for (GVA::RegionOfInterest &roi : regions) {
        // Find tensor generated by reidentification model
        for (GVA::Tensor &tensor : roi.tensors()) {
            if (masterGstElement_->model) {
                if (tensor.model_name().find(masterGstElement_->model) == std::string::npos)
                    continue;
            } else {
                if (tensor.format() != "cosine_distance")
                    continue;
            }

            // embeddings
            std::vector<float> data = tensor.data<float>();
            cv::Mat blob_wrapper(data.size(), 1, CV_32F, data.data());
            embeddings.emplace_back();
            blob_wrapper.copyTo(embeddings.back());
            // tensors
            tensors.push_back(tensor);
            break;
        }
    }

    // Compare embeddings against gallery
    auto ids = gallery->GetIDsByEmbeddings(embeddings);
    if (ids.size() != tensors.size())
        std::logic_error("The number of id received does not match the number of tensors!");

    // Store label for identified objects
    for (size_t i = 0; i < ids.size(); i++) {
        tensors[i].set_string("label", gallery->GetLabelByID(ids[i]));
        if (ids[i] != EmbeddingsGallery::unknown_id) {
            tensors[i].set_int("label_id", ids[i] + 1); // recognized objects starting label_id=1
        } else
            tensors[i].set_int("label_id", EmbeddingsGallery::unknown_id);
    }
}

void Identify::ProcessOutput(GstBuffer *buffer) {
    if (gallery) {
        IdentifyObjects(buffer);
    }
}

GstFlowReturn frame_to_identify(GstGvaIdentify *ovino, GstBuffer *buf) {
    ovino->identifier->ProcessOutput(buf);
    return GST_FLOW_OK;
}

Identify *identifier_new(GstGvaIdentify *ovino) {
    return new Identify(ovino);
}

void identifier_delete(Identify *identifier) {
    delete identifier;
}
